import os
import base64
from volcenginesdkarkruntime import Ark

# 定义方法将指定路径图片转为Base64编码
def encode_image(image_path):
  with open(image_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')


def analyze_two_images(image1_path, image2_path, prompt):
    """调用豆包API分析两张图片，附带提示词"""
    # 转换两张图片为Base64
    img1_base64 = encode_image(image1_path)
    img2_base64 = encode_image(image2_path)

    # content是一个列表，每个元素通过type区分文本或图片
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{img1_base64}"
                    },
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{img2_base64}"
                    },
                },
                # 提示词部分
                {"type": "text", "text": prompt}
            ]
        }
    ]

    try:
        # 从环境变量中读取您的方舟API Key
        client = Ark(
            api_key=os.environ.get("ARK_API_KEY"),
            base_url="https://ark.cn-beijing.volces.com/api/v3",
        )
        # 调用豆包API
        completion = client.chat.completions.create(
            model="doubao-seed-1-6-flash-250828",
            messages = messages
        )
        # 返回模型响应
        return completion.choices[0].message.content
    except Exception as e:
        return f"API调用失败：{str(e)}"


# 在main.py的示例用法部分修改
if __name__ == "__main__":
    from prompts import get_prompt  # 导入提示词工具

    # 图片路径（替换为你的两张图片路径）
    image1 = "images/image_1.png"
    image2 = "images/image_2.png"

    # 获取提示词：可以指定版本和类型
    user_prompt = get_prompt(version="v1")

    # 调用分析函数
    result = analyze_two_images(image1, image2, user_prompt)
    print("分析结果：")
    print(result)